# build_fraud_kg.py 构建图谱代码
import os
import pandas as pd
from neo4j import GraphDatabase

# === Neo4j 配置 ===
URI = "neo4j://127.0.0.1:7687"
USER = "neo4j"
PASSWORD = "lcw123456789"  # ← 你的密码
DATA_DIR = "data"

driver = GraphDatabase.driver(URI, auth=(USER, PASSWORD))

def safe_to_numeric(series):
    return pd.to_numeric(series, errors='coerce')

def clean_role(raw_role):
    """统一角色标签"""
    if pd.isna(raw_role):
        return "其他"
    role_str = str(raw_role)
    if '业务员' in role_str or '代理人' in role_str:
        return '业务员'
    elif '理赔' in role_str:
        return '理赔员'
    elif '核保' in role_str:
        return '核保员'
    elif '管理' in role_str:
        return '管理人员'
    else:
        return '其他'

def main():
    if not os.path.exists(DATA_DIR):
        raise FileNotFoundError(f"请创建 {DATA_DIR}/ 文件夹，并放入5个CSV文件")

    # === 1. 读取数据（自动处理编码）===
    encodings = ['utf-8', 'gbk', 'gb2312']
    dfs = {}
    for file in ['employee.csv', 'policy.csv', 'claim.csv', 'visit.csv', 'institution.csv']:
        for enc in encodings:
            try:
                dfs[file.replace('.csv', '')] = pd.read_csv(
                    f"{DATA_DIR}/{file}", dtype=str, encoding=enc
                )
                print(f"✅ {file} 成功加载（编码: {enc}）")
                break
            except UnicodeDecodeError:
                continue
        else:
            raise ValueError(f"无法解码 {file}，请保存为 UTF-8 或 GBK 格式")

    employees = dfs['employee']
    policies = dfs['policy']
    claims = dfs['claim']
    visits = dfs['visit']
    institutions = dfs['institution']

    # === 2. 转换数值字段 ===
    policies['保费（元）'] = safe_to_numeric(policies['保费（元）'])
    claims['赔付金额（元）'] = safe_to_numeric(claims['赔付金额（元）'])
    visits['总费用（元）'] = safe_to_numeric(visits['总费用（元）'])
    employees['年龄'] = safe_to_numeric(employees['年龄'])

    with driver.session() as session:
        # === 3. 创建约束（加速查询）===
        print("🔄 创建约束...")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (p:Person) REQUIRE p.id IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (pol:Policy) REQUIRE pol.policy_id IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (c:Claim) REQUIRE c.claim_id IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (inst:Institution) REQUIRE inst.inst_id IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (mr:MedicalRecord) REQUIRE mr.record_id IS UNIQUE")

        # === 4. 导入人员（去重 + 角色清洗）===
        print("🔄 导入人员...")
        all_persons = set()

        # 员工
        for _, r in employees.iterrows():
            role = clean_role(r['职位'])
            all_persons.add((str(r['员工ID']), r['姓名'], role))

        # 投保/被保人
        for _, r in policies.iterrows():
            all_persons.add((str(r['投保人ID']), r['投保人姓名'], '客户'))
            all_persons.add((str(r['被保人ID']), r['被保人姓名'], '客户'))

        # 理赔申请人
        for _, r in claims.iterrows():
            all_persons.add((str(r['申请人ID']), r['申请人姓名'], '客户'))

        # 就诊患者
        for _, r in visits.iterrows():
            all_persons.add((str(r['患者ID']), r['患者姓名'], '客户'))

        # 主治医生（作为 Person）
        for _, r in visits.iterrows():
            doc_id = str(r['主治医生'])
            all_persons.add((doc_id, r['主治医生'], '医生'))

        # 理赔员/调查员
        for _, r in claims.iterrows():
            all_persons.add((str(r['受理理赔员ID']), '未知', '理赔员'))
            all_persons.add((str(r['调查人员ID']), '未知', '调查员'))

        for pid, name, role in all_persons:
            session.run("""
                MERGE (p:Person {id: $id})
                SET p.name = COALESCE($name, '未知'), p.role = $role
            """, id=pid, name=name, role=role)

        # === 5. 导入保单 + 关系 ===
        print("🔄 导入保单及关系...")
        for _, r in policies.iterrows():
            session.run("""
                MATCH (applicant:Person {id: $applicant_id})
                MATCH (agent:Person {id: $agent_id})
                CREATE (pol:Policy {
                    policy_id: $policy_id,
                    product: $product,
                    premium: $premium,
                    apply_date: date($apply_date)
                })
                CREATE (applicant)-[:APPLY_FOR]->(pol)
                CREATE (agent)-[:SELLS]->(pol)
            """,
            applicant_id=str(r['投保人ID']),
            agent_id=str(r['销售业务员ID']),
            policy_id=str(r['保单号']),
            product=str(r['险种']),
            premium=r['保费（元）'],
            apply_date=str(r['投保时间']))

        # === 6. 导入理赔 + 关系 ===
        print("🔄 导入理赔及关系...")
        for _, r in claims.iterrows():
            fraud_text = str(r['涉嫌欺诈（标签）']).strip()
            is_fraud = '是' in fraud_text  # 关键：只要包含“是”即为欺诈
            session.run("""
                MATCH (pol:Policy {policy_id: $policy_id})
                MATCH (applicant:Person {id: $applicant_id})
                CREATE (c:Claim {
                    claim_id: $claim_id,
                    amount: $amount,
                    is_fraud: $is_fraud
                })
                CREATE (pol)-[:HAS_CLAIM]->(c)
                CREATE (applicant)-[:APPLIES_FOR]->(c)
            """,
            policy_id=str(r['关联保单号']),
            applicant_id=str(r['申请人ID']),
            claim_id=str(r['理赔案件号']),
            amount=r['赔付金额（元）'],
            is_fraud=is_fraud)

        # === 7. 导入医疗机构 ===
        print("🔄 导入医疗机构...")
        for _, r in institutions.iterrows():
            session.run("""
                MERGE (inst:Institution {inst_id: $inst_id})
                SET inst.name = $name, inst.level = $level
            """,
            inst_id=str(r['医疗机构ID']),
            name=str(r['机构名称']),
            level=str(r['等级']))

        # === 8. 导入就诊记录 + 关系 ===
        print("🔄 导入就诊记录及关系...")
        for _, r in visits.iterrows():
            doc_id = str(r['主治医生'])
            session.run("""
                MATCH (p:Person {id: $patient_id})
                MATCH (inst:Institution {inst_id: $inst_id})
                MERGE (doc:Person {id: $doc_id})
                ON CREATE SET doc.name = $doc_name, doc.role = '医生'
                CREATE (mr:MedicalRecord {
                    record_id: $record_id,
                    diagnosis: $diagnosis,
                    cost: $cost
                })
                CREATE (p)-[:VISITS]->(mr)
                CREATE (mr)-[:TREATED_AT]->(inst)
                CREATE (doc)-[:TREATS]->(mr)
            """,
            patient_id=str(r['患者ID']),
            inst_id=str(r['医疗机构ID']),
            doc_id=doc_id,
            doc_name=str(r['主治医生']),
            record_id=str(r['就诊ID']),
            diagnosis=str(r['诊断结果']),
            cost=r['总费用（元）'])

    print("🎉 知识图谱构建完成！")

if __name__ == "__main__":
    main()